Skip to content

[Quantization] Fix ModelOpt pre-quantized loading #14188

Open
yzhautouskay wants to merge 2 commits into
huggingface:mainfrom
yzhautouskay:yzhautouskay/modelopt_prequant_fix
Open

[Quantization] Fix ModelOpt pre-quantized loading #14188
yzhautouskay wants to merge 2 commits into
huggingface:mainfrom
yzhautouskay:yzhautouskay/modelopt_prequant_fix

Conversation

@yzhautouskay

@yzhautouskay yzhautouskay commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes loading of pre-quantized ModelOpt models via from_pretrained(..., device_map=...),
which currently crashes with:

NotImplementedError: Cannot copy out of meta tensor; no data! ...

Note: this is required for the pre-quantized path to work at all with HF checkpointing enabled —
no combination of device_map / .to(device) load options works around it without this change.

Root cause

For pre-quantized models, NVIDIAModelOptQuantizer.check_if_quantized_param returned True
for every tensor. That routes the quantizer buffers ModelOpt recreates on restore
(_amax, _scale) through create_quantized_param, which assigns them into
module._parameters. Those buffers are then never materialized by the low-memory loader and
stay on the meta device, so the later device move fails.

Fix

  • check_if_quantized_param: for pre-quantized models, only claim real parameters
    (return tensor_name in module._parameters); buffers fall back to the standard loader, which
    materializes _amax/_scale on their target device.
  • create_quantized_param: guard the pre-quantized branch so a non-parameter reaching it
    raises clearly instead of silently mis-assigning.

Regression test

Added test_prequantized_serialization_with_device_map, which:

  • Enables enable_huggingface_checkpointing() before quantization, save, and reload.
  • Verifies modelopt_state.pth is produced.
  • Reloads the pre-quantized checkpoint with device_map.
  • Confirms ModelOpt _amax/_scale buffers are restored.
  • Asserts no parameters or buffers remain on meta.

Validation:

RUN_NIGHTLY=1 python -m pytest \
  tests/quantization/modelopt/test_modelopt.py::SanaTransformerFP8WeightsTest::test_prequantized_serialization_with_device_map \
  -v -s
  • Without this fix: fails with NotImplementedError: Cannot copy out of meta tensor; no data!
  • With this fix: passes

Environment

Single ModelOpt version on both export and load (ModelOpt state is version-sensitive, so
producer and consumer must match):

  • diffusers 0.40.0.dev0 (this branch)
  • nvidia-modelopt 0.44.0 (quantize + load)
  • torch 2.9.0a0+145a3a7bda.nv25.10 (NVIDIA NGC container build)

Who can review?

@sayakpaul

@github-actions github-actions Bot added quantization size/S PR with diff < 50 LOC labels Jul 14, 2026
@github-actions

Copy link
Copy Markdown
Contributor

Hi @yzhautouskay, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant